# -*- coding: utf-8 -*-
"""
@Time ： 2024/1/23 21:10
@Auth ： 狗头老道
@File ：MaskRasterBatch.py
@IDE ：PyCharm
@Purpose ：基于GDAL批量裁剪tif格式栅格
"""
# -*- coding: utf-8 -*-

import time
import os
from osgeo import gdal
from PyQt5 import QtWidgets


# from widgetAndDialog import ProgressBar, pyqtbar

def clip_batch(in_path, out_path, in_shape, prefix="" ):
    """
    :param in_path: 需要裁剪栅格文件的文件夹
    :param out_path: 输出文件夹
    :param in_shape: shp文件
    :param prefix: 裁剪文件的前缀
    :return:    """

    # 定义shp文件的目录+名称
    files = os.listdir(in_path)
    # 打开需要裁剪的文件夹,将所有文件以列表的形式列出

    for file in files:
        if file[-4:] == '.tif':
            # 判断文件是否为.tif结尾
            filename = os.path.join(in_path, file)
            # 确定找到的文件名
            in_raster = gdal.Open(filename)
            out_raster = os.path.join(out_path, prefix + file)

            # ------------------------------------影像基础信息------------------------------------
            ds_width = in_raster.RasterXSize  # 获取数据宽度
            ds_height = in_raster.RasterYSize  # 获取数据高度
            ds_geo = in_raster.GetGeoTransform()  # 获取仿射地理变换参数
            ds_prj = in_raster.GetProjection()  # 获取投影信息
            ds_nodata = in_raster.GetRasterBand(1).GetNoDataValue()  # 获取影像的nodata值

            ds = gdal.Warp(out_raster, in_raster, format='GTiff',
                           cutlineDSName=in_shape,
                           cropToCutline=True,
                           cutlineWhere=None, dstNodata=ds_nodata)
            ds = None
            # 关闭处理空间，释放内存


def mask_to_raster(in_raster_file, out_path, in_shape,  prefix=""):
    if in_raster_file[-4:] == '.tif':# 判断文件是否为.tif结尾
        # 确定找到的文件名
        # 获取基本文件名
        filename = os.path.basename(in_raster_file)
        in_raster = gdal.Open(in_raster_file)

        out_raster = os.path.join(out_path, prefix + filename)

        # ------------------------------------影像基础信息------------------------------------
        ds_width = in_raster.RasterXSize  # 获取数据宽度
        ds_height = in_raster.RasterYSize  # 获取数据高度
        ds_geo = in_raster.GetGeoTransform()  # 获取仿射地理变换参数
        ds_prj = in_raster.GetProjection()  # 获取投影信息
        ds_nodata = in_raster.GetRasterBand(1).GetNoDataValue()  # 获取影像的nodata值

        ds = gdal.Warp(out_raster, in_raster, format='GTiff',
                       cutlineDSName=in_shape,
                       cropToCutline=True,
                       cutlineWhere=None, dstNodata=ds_nodata)
        ds = None


if __name__ == "__main__":
    # 直接执行函数
    start = time.perf_counter()  # 开始时间
    in_shape = r"D:\test\阿克苏土地利用变化模拟\mask.shp"  # 矢量范围
    in_path = r"D:\test\阿克苏土地利用变化模拟"  # 输入栅格路径
    out_path = r"D:\test\阿克苏土地利用变化模拟\mask"  # 输出栅格路径
    clip_batch(in_path, out_path, in_shape, "clip_")
    end = time.perf_counter()  # 结束时间
    print('finish')
    print('Running time: %s Seconds' % (end - start))
    # 展示程序运行时间
